# @Time : 2022/11/18 15:18
# @Author : xiashuobad
# @File : color_clustering_and_matching.py
# @Software: PyCharm
"""
/**
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
from pathlib import Path

from PIL import Image
import numpy as np
from sklearn.cluster import KMeans


def centroid_histogram(clt):
    numLabels = np.arange(0, len(np.unique(clt.labels_)) + 1)
    (hist, _) = np.histogram(clt.labels_, bins=numLabels)
    hist = hist.astype("float")
    hist /= hist.sum()
    return hist


def color_clustering(img_path, num_class):
    image = Image.open(img_path)
    ori_size = image.size
    image = image.resize((512, 320))
    image_k = np.array(image)
    if image_k.shape[-1] > 3:
        image_k = image_k[..., :3]
    image_k = image_k.reshape((-1, 3))
    clt = KMeans(n_clusters=num_class)
    clt.fit(image_k)
    hist: np.ndarray = centroid_histogram(clt)
    list_rgb = clt.cluster_centers_.astype("int")
    list_rgb = list_rgb[hist.argsort()]  # 按颜色占比从小到大排序
    image = image.resize(ori_size)
    return list_rgb, image


def class_histogram(mask):
    if isinstance(mask, str):
        mask = Image.open(mask)
    mask = np.array(mask)
    assert len(mask.shape) == 2
    classes, counts = np.unique(mask, return_counts=True)  # 利用np的unique函数得到所有灰度值以及各个灰度值的计数
    hist = counts / counts.sum()  # 等比例变形，以最小计数的灰度值的比例为1，这样方便查看各个灰度值计数之间的倍数关系
    return classes, hist


def color_matching(mask_path, color_map, save_path=None):
    mask = Image.open(mask_path).convert("L")  # type:# Image.Image
    classes, hist = class_histogram(mask)
    color_map = color_map[hist.argsort().argsort()]
    color_map = color_map.reshape((-1,)).tolist()
    assert len(color_map) >= len(classes) * 3
    mask = mask.convert('L')
    mask.putpalette(color_map)
    save_path and mask.save(save_path)
    return mask


def main(image_ori_path, image_color_path, image_target_path, num_class, save_path=None):
    image_ori: Image.Image = Image.open(image_ori_path)
    width, height = image_ori.size
    image_showcase = Image.new(image_ori.mode, (width * 2, height * 2), color=(255, 255, 255))
    image_showcase.paste(image_ori, (0, 0))
    list_rgb, image_color = color_clustering(image_color_path, num_class)  # type:# np.ndarray
    image_color = resize_img(image_color, (width, height), (255, 255, 255))
    image_showcase.paste(image_color, (width, 0))
    image_target = color_matching(image_target_path, list_rgb, save_path)
    image_showcase.paste(image_target, (width // 2, height))
    image_showcase.show()
    save_path and image_showcase.save(save_path.replace(suffix:=Path(save_path).suffix,f'-showcase{suffix}'))


def resize_img(img: (str, Image.Image), target_size: tuple, pad=None):
    """
    使用Image对图片进行risize
    Args:
        img:
        target_size:(w,h)

    Returns:

    """
    if isinstance(img, str):
        img: Image.Image = Image.open(img)
    assert isinstance(img, Image.Image)
    ori_w, ori_h = img.size
    new_w, new_h = target_size
    if pad is None:  # 不填充，直接拉伸
        new_img = img.resize(target_size, Image.NEAREST)
    else:  # 保持横纵比，其余部分填充
        assert isinstance(pad, (int, tuple)), 'pad参数应该为整型数值或者rgb元组，表示要填充的像素值'
        new_img = Image.new(img.mode, target_size, pad)
        scale_rate_w, scale_rate_h = new_w / ori_w, new_h / ori_h
        if scale_rate_w < scale_rate_h:
            img = img.resize((new_w, cur_h := int(ori_h * scale_rate_w)))
            new_img.paste(img, (0, (new_h - cur_h) // 2))
        else:
            img = img.resize((cur_w := int(ori_w * scale_rate_h), new_h))
            new_img.paste(img, ((new_w - cur_w) // 2, 0))

    return new_img


if __name__ == '__main__':
    image_color_path = r'images\color\color4.jpg'
    image_target_path = r'images\target\411381dengzhou-1.tif'
    image_ori_path = r'images\target\411381dengzhou-1-ori.tif'
    save_path = r'images\target/411381dengzhou-1-color4.png'
    main(image_ori_path, image_color_path, image_target_path, 8, save_path)
